package com.liyunc.demo.detect.service;

import com.liyunc.demo.detect.recognition.PredictRlt;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.opencv_core;
import org.bytedeco.javacpp.opencv_objdetect;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.OpenCVFrameConverter;

import java.io.File;
import java.util.Map;

import static org.bytedeco.javacpp.opencv_core.FONT_HERSHEY_PLAIN;
import static org.bytedeco.javacpp.opencv_imgproc.CV_AA;
import static org.bytedeco.javacpp.opencv_imgproc.CV_BGR2GRAY;
import static org.bytedeco.javacpp.opencv_imgproc.cvtColor;
import static org.bytedeco.javacpp.opencv_imgproc.putText;
import static org.bytedeco.javacpp.opencv_imgproc.rectangle;

/**
 * 检测人脸并认证的服务.
 * 在画面上标识出检测结果.
 */
@Slf4j
public class DetectAndRecognizeService implements DetectService {

    private static final double FONT_SCALE = 1.5;

    //每一帧原始图片
    private opencv_core.Mat grabbedImage = null;

    // 原始图片对应的灰度图片对象
    private opencv_core.Mat grayImage = null;

    //分类器
    private opencv_objdetect.CascadeClassifier classifier;

    //转换器
    private final OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();

    // 模型文件地址
    private final String detectModelFilePath;

    // 处理每一帧的服务
    private final RecognizeService recognizeService;

    // 为了显示的时候更加友好，给每个分类对应一个名称
    private final Map<Integer, String> kindNameMap;

    /**
     * 构造方法，指定模型文件的下载地址.
     *
     * @param detectModelFilePath    人脸检测模型地址
     * @param recognizeModelFilePath 推测模型文件地址
     * @param kindNameMap            类别
     */
    public DetectAndRecognizeService(String detectModelFilePath, String recognizeModelFilePath, Map<Integer, String> kindNameMap) {
        this.detectModelFilePath = detectModelFilePath;
        this.recognizeService = new RecognizeService(recognizeModelFilePath);
        this.kindNameMap = kindNameMap;
    }

    @Override
    public void init() {
        // 打开模型文件
        File file = new File(detectModelFilePath);

        // 模型文件下载后的完整地址
        String classifierName = file.getAbsolutePath();

        //根据模型文件实例化分类器
        classifier = new opencv_objdetect.CascadeClassifier(classifierName);
    }

    @Override
    public Frame convert(Frame frame) {
        // 由帧转为Mat
        grabbedImage = converter.convert(frame);

        // 灰度Mat，用于检测
        if (null == grayImage) {
            grayImage = DetectService.buildGrayImage(grabbedImage);
        }

        // 进行人脸识别
        return detectAndRecognize(classifier, converter, frame, grabbedImage,
                grayImage, recognizeService, kindNameMap);
    }

    @Override
    public void releaseResource() {
        log.info("releaseResource");
        if (null != grabbedImage) {
            grabbedImage.release();
        }
        if (null != grayImage) {
            grayImage.release();
        }
        if (null != classifier) {
            classifier.close();
        }
    }

    /**
     * 检测图片，将检测结果用矩形标注在原始图片上.
     *
     * @param classifier       分类器
     * @param converter        转换器
     * @param frame            原始视频帧
     * @param grabbedImage     原始帧对应的Mat
     * @param grayImage        存放灰度图片的Mat
     * @param recognizeService 认证服务
     * @param kindNameMap      每个分类对应的编号
     * @return 标注了结果的视频帧
     */
    private Frame detectAndRecognize(opencv_objdetect.CascadeClassifier classifier,
                                     OpenCVFrameConverter.ToMat converter,
                                     Frame frame,
                                     opencv_core.Mat grabbedImage,
                                     opencv_core.Mat grayImage,
                                     RecognizeService recognizeService,
                                     Map<Integer, String> kindNameMap) {
        // 当前图片转为灰度图片
        cvtColor(grabbedImage, grayImage, CV_BGR2GRAY);

        // 存放检测结果的容器
        opencv_core.RectVector objects = new opencv_core.RectVector();

        // 开始检测
        classifier.detectMultiScale(grayImage, objects);

        // 检测结果总数
        long total = objects.size();

        // 如果没有检测到结果，就用原始帧返回
        if (total < 1) {
            log.info("detect result none ");
            return frame;
        }

        PredictRlt predictRlt;
        int posX;
        int posY;
        int label;
        double confidence;
        String content;

        // 如果有检测结果，就根据结果的数据构造矩形框，画在原图上
        for (long i = 0; i < total; i++) {
            opencv_core.Rect rect = objects.get(i);

            // 核心代码，把检测的人脸拿去识别
            predictRlt = recognizeService.predict(new opencv_core.Mat(grayImage, rect));

            // 如果返回为空，表现出现异常，继续执行下一个
            if (null == predictRlt) {
                log.info("recognizeService abnormal predict");
                continue;
            }

            // 分类的编号（训练时只有1和2，这里只有有三个值，1和2与训练的分类一致，还有个-1表示没有匹配上）
            label = predictRlt.getLabel();
            // 与模型中的分类的距离，值越小表示相似度越高
            confidence = predictRlt.getConfidence();

            // 得到分类编号后，从map中取得名字，用来显示
            if (kindNameMap.containsKey(label)) {
                content = String.format("%s, confidence : %.4f",
                        kindNameMap.get(label), confidence);
            } else {
                content = "unknown[" + predictRlt.getLabel() + "]";
            }
            log.info("recognizeService predict result: {}", content);

            // 人脸的位置信息
            int rectX = rect.x();
            int rectY = rect.y();
            int rectW = rect.width();
            int rectH = rect.height();

            // 在人脸上画矩形
            rectangle(grabbedImage,
                    new opencv_core.Point(rectX, rectY),
                    new opencv_core.Point(rectX + rectW, rectY + rectH),
                    opencv_core.Scalar.RED, 1, CV_AA, 0);

            final int position = 10;
            posX = Math.max(rect.tl().x() - position, 0);
            posY = Math.max(rect.tl().y() - position, 0);

            putText(grabbedImage,
                    content,
                    new opencv_core.Point(posX, posY),
                    FONT_HERSHEY_PLAIN,
                    FONT_SCALE,
                    opencv_core.Scalar.GREEN);
        }

        // 释放检测资源
        objects.close();

        // 将标注过的图片转为帧，返回
        return converter.convert(grabbedImage);
    }
}
